#!/usr/bin/python

# hostbuilder - Tools for DNS zonefile administration
# Copyright (c) 2010 Crown copyright
# 
# This file is part of hostbuilder.
# 
# hostbuilder is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# hostbuilder is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with hostbuilder.  If not, see <http://www.gnu.org/licenses/>.


import sys
import re

import dns.name

import socket
socket.setdefaulttimeout(10)


def query_nameservers(zone):
    import dns.resolver
    nameservers = []
    ans = dns.resolver.query(zone, rdtype='NS')
    for a in ans:
        t = a.target.derelativize(zone).to_text(omit_final_dot=True)
        if t.endswith('.'):
            t = t[:-1]
        nameservers.append(t)
    return nameservers


def load_zone(zone, nameserver=None):
    """Load a zone from a zone transfer"""

    import dns.zone
    import dns.query
    if nameserver is None:
        nameservers = query_nameservers(zone)

        while True:
            # iterate over all nameservers until one works
            ns = nameservers.pop(0)
            try:
                return load_zone(zone, ns)
            except socket.timeout:
                if not nameservers:
                    raise

    return dns.zone.from_xfr(dns.query.xfr(nameserver, zone))


def print_dynamic_records(z):
    """Print all of the dynamic records in the zone.
    
    Ddynamic records are identified by the presence of a 32-character
    hexadecimal hash in a corresponding TXT record.
    
    """
    print "; Dynamic records"
    print "$TTL 300"
    for name, ttl, rdata in z.iterate_rdatas('A'):
        txt = z.get_rdataset(name, 'TXT')
        if txt:
            for t in txt:
                if re.match(r'^[a-f0-9]{32}', t.strings[0]):
                    n = name.derelativize(z.origin).to_text()
                    print "%s    IN    A      %s" % (n.ljust(24), rdata.address.ljust(24))
                    print "\t\t\t\tIN\tTXT\t%s\n" % t.strings[0]


def get_leases(z):
    """Extract a mapping of all DHCP leases in the zone"""
    leases = {}
    for name, ttl, rdata in z.iterate_rdatas('A'):
        txt = z.get_rdataset(name, 'TXT')
        if txt:
            for t in txt:
                if re.match(r'^[a-f0-9]{32}', t.strings[0]):
                    n = name.derelativize(z.origin).to_text()
                    leases[rdata.address] = n
    return leases


def print_reverse_dynamic_records(z, origin):
    origin = dns.name.Name(origin.split('.'))
    leases = get_leases(z)
    print "; Dynamic records"
    print "$TTL 300"
    for address, name in leases.items():
        labels = address.split('.')
        labels.reverse()
        labels += ['in-addr', 'arpa']
        a = dns.name.Name(labels)
        if a.is_subdomain(origin):
            print "%s    IN    PTR    %s" % (a.relativize(origin).to_text().ljust(24), name.ljust(24))



if __name__ == '__main__':
    from optparse import OptionParser
    parser = OptionParser()
    parser.add_option('-H', '--host', help="Query a specific DNS server for dynamic records")
    parser.add_option('-r', '--reverse', help="Use this origin and output reverse DNS zone file")

    options, args = parser.parse_args()

    if not args:
        parser.error("You must provide a zone to query.")

    zname = args[0]
    try:
        zone = load_zone(zname, nameserver=options.host)
    except socket.timeout:
        print >>sys.stderr, "No nameservers responded to AXFR query."
        sys.exit(1)

    if options.reverse:
        print_reverse_dynamic_records(zone, options.reverse)
    else:
        print_dynamic_records(zone)
